package bacom;

import java.io.*;
import java.util.*;
import static java.lang.Math.*;


import umontreal.iro.lecuyer.probdist.ChiSquareNoncentralDist;
import umontreal.iro.lecuyer.probdist.ChiSquareDist;
import umontreal.iro.lecuyer.probdist.NormalDist;


/* class Noticcor*/

public class Noticcor {

    String[] chrInfo;
    double[] cnXa;
    double[] cnXb;
    double[] copyNumber;
    double[] correctedCopyNumber;
    double[] segMean;
    double[] correctedSegMean;
    boolean[] isGenotypeAB;
    int numLoci;
    double alpha;
    ArrayList<String> deletionSegChrID;
    ArrayList<Integer> deletionSegStartLocation;
    ArrayList<Integer> deletionSegEndLocation;
    ArrayList<Double> deletionSegProbHemiDeletion;
    ArrayList<Double> deletionSegProbHomoDeletion;
    ArrayList<Double> deletionSegNormalTissueFraction;


    public Noticcor(String[] chrID, double[] copyNumber, double[] alleleA, double[] alleleB,
		    boolean[] isGenotypeAB, double[] segMean) {

	this.numLoci = chrID.length;
	this.chrInfo = chrID;
	this.cnXa = alleleA;
	this.cnXb = alleleB;
	this.copyNumber = copyNumber;
	this.isGenotypeAB = isGenotypeAB;
	this.segMean = segMean;
	correctedSegMean = new double[numLoci];
	correctedCopyNumber = new double[numLoci];

	alpha = 0;

	for (int i = 0; i < numLoci; i ++) {
	    segMean[i] = segMean[i] - 2;   // Let normal segments have mean 0
	}


	deletionSegChrID = new ArrayList<String>();
	deletionSegStartLocation = new ArrayList<Integer>();
	deletionSegEndLocation = new ArrayList<Integer>();
	deletionSegProbHemiDeletion = new ArrayList<Double>();
	deletionSegProbHomoDeletion = new ArrayList<Double>();
	deletionSegNormalTissueFraction = new ArrayList<Double>();


	normalization();
	
    }


    private void normalization() {

	double avgNormalCN = 0;
	double avgNormalA = 0;
	double avgNormalB = 0;
	double numNormalLoci = 0;
	double numNormalLociAB = 0;

	final double NORMAL_SEG_MARGIN = 0.1;

	/* Global normaliztion */

	for (int i = 0; i < numLoci; i ++) {

	    if (abs(segMean[i]) <= NORMAL_SEG_MARGIN) {  
		avgNormalCN += copyNumber[i];
		numNormalLoci ++;
		if (isGenotypeAB[i]) {
		    avgNormalA += cnXa[i];
		    avgNormalB += cnXb[i];
		    numNormalLociAB ++;
		}
	    }

	}

	avgNormalCN = avgNormalCN / numNormalLoci;
	avgNormalA = avgNormalA / numNormalLociAB;
	avgNormalB = avgNormalB / numNormalLociAB;


	for (int i = 0; i < numLoci; i ++) {

	    copyNumber[i] = copyNumber[i] / avgNormalCN * 2;
	    cnXa[i] = cnXa[i] / avgNormalA;
	    cnXb[i] = cnXb[i] / avgNormalB;

	}


	/* Local normaliztion */
	
	int index = 0;
	while (index < numLoci) {
	    String chr = chrInfo[index];
	    int chrStartPos = index;
	    avgNormalCN = 0;
	    avgNormalA = 0;
	    avgNormalB = 0;
	    numNormalLoci = 0;
	    numNormalLociAB = 0;
	    if (abs(segMean[index]) <= Double.MIN_VALUE) {
		avgNormalCN += copyNumber[index];
		numNormalLoci ++;
		if (isGenotypeAB[index]) {
		    avgNormalA += cnXa[index];
		    avgNormalB += cnXb[index];
		    numNormalLociAB ++;
		}

	    }

	    index ++;		

	    while (chr.equals(chrInfo[index]) && index < numLoci) {

		if (abs(segMean[index]) <= Double.MIN_VALUE) {
		    avgNormalCN += copyNumber[index];
		    numNormalLoci ++;
		    if (isGenotypeAB[index]) {
			avgNormalA += cnXa[index];
			avgNormalB += cnXb[index];
			numNormalLociAB ++;
		    }
		}
		index ++;		
		if (index >= numLoci) {
		    break;
		}
	    }

	    if (numNormalLociAB > 500) {   /* if the number of normal loci with genotyp AB is greater
					      than 500, then perform normalization.*/

		avgNormalCN = avgNormalCN / numNormalLoci;
		avgNormalA = avgNormalA / numNormalLociAB;
		avgNormalB = avgNormalB / numNormalLociAB;

		for (int i = chrStartPos; i < index; i++) {
		    copyNumber[i] = copyNumber[i] / avgNormalCN * 2;
		    cnXa[i] = cnXa[i] / avgNormalA;
		    cnXb[i] = cnXb[i] / avgNormalB;
		}
	    }

	    
	}

    }




    public void correctNTC() {
	int startPos = 0;
	int endPos = -1;
	int m = 0;
	ArrayList<Double> segCN = new ArrayList<Double>(10000);
	ArrayList<Double> segXa = new ArrayList<Double>(10000);
	ArrayList<Double> segXb = new ArrayList<Double>(10000);

	double muAB;
	Double[] posteriorProb;
	double rhoAB = 0;
	String atChr = null;
	
	rhoAB = robustRho();

	alpha = 0;
	double sumSegSize = 0;

	for (int i = 1; i < numLoci; i++) {

	    segCN.clear();
	    segXa.clear();
	    segXb.clear();

	    if (segMean[i-1] < 0 && 
		(abs(segMean[i] - segMean[i-1]) > 0.000001 || i == (numLoci-1))) {
		

		if (i == (numLoci-1)) {
		    endPos = i;
		} else {
		    endPos = i - 1;
		}

		for (int j = startPos; j <= endPos; j++) {
		    segCN.add(copyNumber[j]);
		    if (isGenotypeAB[j]) {
			segXa.add(cnXa[j]);
			segXb.add(cnXb[j]);
		    }
		}
		

		if (segXa.size() >= 5 && (!chrInfo[endPos].equals("MT"))
		    && (!chrInfo[endPos].equals("X")) && (!chrInfo[endPos].equals("Y")) ) {
		    muAB = robustMu(segCN);
		    posteriorProb = bayesianPosterior(segCN, segXa, segXb, muAB, rhoAB, 0.25);
		    atChr = chrInfo[endPos];

		    if (!Double.isNaN(posteriorProb[2])) {
			alpha = alpha + (1 - posteriorProb[2]) * posteriorProb[3];
			sumSegSize = sumSegSize + posteriorProb[3];
		    }

		    
		    deletionSegChrID.add(atChr);
		    deletionSegStartLocation.add(startPos);
		    deletionSegEndLocation.add(endPos);
		    deletionSegProbHemiDeletion.add(posteriorProb[0]);
		    deletionSegProbHomoDeletion.add(posteriorProb[1]);
		    deletionSegNormalTissueFraction.add(1 - posteriorProb[2]);

		}
		
	    } 
	    
	    if (segMean[i] < 0 && 
		(abs(segMean[i] - segMean[i-1]) > Double.MIN_VALUE)) {

		startPos = i;
	    }
	}

        alpha = alpha / sumSegSize;


    }



    private double robustMu(ArrayList<Double> segCN) {

	ArrayList<Double> sumCNaCNb = new ArrayList<Double>(segCN.size());
	double sigma = 0;
	NormalDist gaussian = new NormalDist();
	double segSize = 0;
	double cutoff = 0;
	double cutoff1 = 0;
	double cutoff2 = 0;

	for (int i = 0; i< segCN.size(); i++) {  // Remove outliers before compute mu
	    if (segCN.get(i) >= -50 && segCN.get(i) <= 50) {
		sumCNaCNb.add(segCN.get(i));
	    }
	}

	segSize = sumCNaCNb.size();
	
	sigma = std(sumCNaCNb);
	double areaRemoved = 0.5/segSize;

	cutoff = abs(gaussian.inverseF(areaRemoved));
	
	cutoff1 =  -cutoff * sigma + mean(sumCNaCNb);
	cutoff2 = cutoff * sigma + mean(sumCNaCNb);

	while (true) {


	    Iterator<Double> iter = sumCNaCNb.iterator();
	    
	    int numRemoved = 0;
	    while (iter.hasNext()) {

		double cn = iter.next();
		if (cn < cutoff1 || cn > cutoff2) {
		    numRemoved++;
		    iter.remove();
		}
	    }

	    if (numRemoved == 0) {
		break;
	    }

	    
	    sigma = std(sumCNaCNb);
	    sigma = sigma / sqrt( (1 - 2*areaRemoved - 2*cutoff*exp(-pow(cutoff,2)/2)/sqrt(2*PI)) / (1 - 2*areaRemoved) );

	    segSize = sumCNaCNb.size();
	    areaRemoved = 0.5/segSize;

	    cutoff = abs(gaussian.inverseF(areaRemoved));

	    cutoff1 =  -cutoff * sigma + mean(sumCNaCNb);
	    cutoff2 = cutoff * sigma + mean(sumCNaCNb);

	    
 	}

	return mean(sumCNaCNb);
    }




    private double robustRho() {

	ArrayList<Double> sumCNaCNb = new ArrayList<Double>(cnXa.length);
	ArrayList<Double> sXa = new ArrayList<Double>(cnXa.length);
	ArrayList<Double> sXb = new ArrayList<Double>(cnXa.length);
	double sigma = 0;
	NormalDist gaussian = new NormalDist();
	double segSize = 0;
	double cutoff = 0;
	double cutoff1 = 0;
	double cutoff2 = 0;

	for (int i = 0; i< cnXa.length; i++) {  // Remove outliers before compute mu
	    if (copyNumber[i] >= -50 && copyNumber[i] <= 50 && isGenotypeAB[i] && abs(segMean[i]) <= 0.001) {
		sumCNaCNb.add(copyNumber[i]);
		sXa.add(cnXa[i]);
		sXb.add(cnXb[i]);
	    }
	}

	segSize = (double) sumCNaCNb.size();
	
	sigma = std(sumCNaCNb);
	double areaRemoved = 0.5/segSize;

	cutoff = abs(gaussian.inverseF(areaRemoved));
	
	cutoff1 =  -cutoff * sigma + mean(sumCNaCNb);
	cutoff2 = cutoff * sigma + mean(sumCNaCNb);

	while (true) {


	    Iterator<Double> iter = sumCNaCNb.iterator();
	    Iterator<Double> iterXa = sXa.iterator();
	    Iterator<Double> iterXb = sXb.iterator();

	    int numRemoved = 0;
	    while (iter.hasNext()) {

		double cn = iter.next();
		iterXa.next();
		iterXb.next();
		if (cn < cutoff1 || cn > cutoff2) {
		    numRemoved++;
		    iter.remove();
		    iterXa.remove();
		    iterXb.remove();
		}
	    }

	    if (numRemoved == 0) {
		break;
	    }

	    
	    sigma = std(sumCNaCNb);
	    sigma = sigma / sqrt( (1 - 2*areaRemoved - 2*cutoff*exp(-pow(cutoff,2)/2)/sqrt(2*PI)) / (1 - 2*areaRemoved) );

	    segSize = sumCNaCNb.size();
	    areaRemoved = 0.5/segSize;

	    cutoff = abs(gaussian.inverseF(areaRemoved));

	    cutoff1 =  -cutoff * sigma + mean(sumCNaCNb);
	    cutoff2 = cutoff * sigma + mean(sumCNaCNb);
	    
 	}

	double corr = 0;
	double meanXa = mean(sXa);
	double meanXb = mean(sXb);

	for (int i = 0; i < sXa.size(); i++) {
	    corr = corr + (sXa.get(i) - meanXa) * (sXb.get(i) - meanXb);
	}

	corr = (corr / (sXa.size() - 1)) / (std(sXa) * std(sXb));

	return corr;
    }






    private double mean(ArrayList<Double> x) {
	double sum = 0;
	double length = x.size();
	for (int i = 0; i < x.size(); i++) {
	    sum = sum + x.get(i);
	}
	return (sum / length);
    }


    private double std(ArrayList<Double> x) {
	double sum = 0;
	double meanX = mean(x);
	double length = x.size();
	for (int i = 0; i < x.size(); i++) {
	    sum = sum + pow((x.get(i)-meanX), 2);
	}
	return (sqrt(sum / (length-1)));
    }


    private double std(ArrayList<Double> x, double meanX) {
	double sum = 0;
	double length = x.size();
	for (int i = 0; i < x.size(); i++) {
	    sum = sum + pow((x.get(i)-meanX), 2);
	}
	return (sqrt(sum / (length)));
    }








    private Double[] bayesianPosterior(ArrayList<Double> segCN, ArrayList<Double> segXa, ArrayList<Double> segXb, 
				      double muAB, double rhoAB, double pThreshold) {

	ArrayList<Double> sumXaXb = new ArrayList<Double>(segXa.size());
	ArrayList<Double> sXa = new ArrayList<Double>(segXa.size());
	ArrayList<Double> sXb = new ArrayList<Double>(segXb.size());
	int segSize = segXa.size();

        for (int i = 0; i < segSize; i++) {
	    sumXaXb.add(segCN.get(i));
	    sXa.add(segXa.get(i));
	    sXb.add(segXb.get(i));
	}
	
	// *******************************
	// To remove outliers in the data
	// *******************************
	double sigma = std(sumXaXb, muAB);
	double areaRemoved = 0.5/segSize;
	NormalDist gaussian = new NormalDist();
	double cutoff = 0;
	double cutoff1 = 0;
	double cutoff2 = 0;

	cutoff = abs(gaussian.inverseF(areaRemoved));
	
	cutoff1 =  -cutoff * sigma + muAB;
	cutoff2 = cutoff * sigma + muAB;

	while (true) {

	    Iterator<Double> iter = sumXaXb.iterator();
	    Iterator<Double> iterXa = sXa.iterator();
	    Iterator<Double> iterXb = sXb.iterator();

	    int numRemoved = 0;
	    while (iter.hasNext()) {

		double cn = iter.next();
		iterXa.next();
		iterXb.next();
		if (cn < cutoff1 || cn > cutoff2) {
		
		    numRemoved++;
		    iter.remove();
		    iterXa.remove();
		    iterXb.remove();

		}
	    }

	    if (numRemoved == 0) {
		break;
	    }

	    sigma = std(sumXaXb, muAB);
	    sigma = sigma / sqrt( (1 - 2*areaRemoved - 2*cutoff*exp(-pow(cutoff,2)/2)/sqrt(2*PI)) / (1 - 2*areaRemoved) );


	    segSize = sumXaXb.size();
	    areaRemoved = 0.5/segSize;
	    cutoff = abs(gaussian.inverseF(areaRemoved));
	    cutoff1 =  -cutoff * sigma + muAB;
	    cutoff2 = cutoff * sigma + muAB;


 	}

	// *******************************

	segSize = sumXaXb.size();
	areaRemoved = 0.5/segSize;

	sigma = std(sumXaXb, muAB);
 	
	sigma = sigma / sqrt( (1 - 2*areaRemoved - 2*cutoff*exp(-pow(cutoff,2)/2)/sqrt(2*PI)) / (1 - 2*areaRemoved) );
	sigma = sigma * sqrt((1 - 1*rhoAB) / (1 + 1*rhoAB)); // corrected sigma
	
	double lambda = (2 - muAB) / sigma;
	lambda = pow(lambda, 2);
	lambda = lambda * segSize;

	// ***********************************
	// New method, using the whole segment
	// ***********************************

	double score = 0;
	for (int i = 0; i < segSize; i++) {
	    score = score + pow(sXa.get(i) - sXb.get(i), 2) / pow(sigma, 2);
	}

	ChiSquareDist chiSquare = new ChiSquareDist(segSize-1);

	ChiSquareNoncentralDist nx2 = new ChiSquareNoncentralDist(segSize - 1, lambda);

// 	System.out.println("\n----------------------------------");
// 	System.out.println("  Chi square parameters: df = " + (segSize - 1));
// 	System.out.println("  Noncentral chi square parameters: lambda = " + lambda);
// 	System.out.println("  score = " + score);
// 	System.out.println("----------------------------------\n\n");


	Double L1 = nx2.density(score);
	Double L2 = chiSquare.density(score);
	
	Double pHemi = new Double(0);
	Double pHomo = new Double(0);

	if (Double.isNaN(L1) && !Double.isNaN(L2)) {
	    pHemi = 0.0;
	    pHomo = 1.0;

	} 
	
	if (!Double.isNaN(L1) && Double.isNaN(L2)) {
	    pHemi = 1.0;
	    pHomo = 0.0;

	} 
	
	if (!Double.isNaN(L1) && !Double.isNaN(L2)) {
	    if (L1 == 0 && L2 != 0) {
		pHemi = 0.0;
		pHomo = 1.0;
	    } 
	    if (L2 == 0 && L1 != 0) {
		pHemi = 1.0;
		pHomo = 0.0;
	    }
	    if (L2 != 0 && L1 != 0) {
		pHemi = L1 / (L1 + L2);
		pHomo = L2 / (L1 + L2);
	    }
	    if (L2 == 0 && L1 == 0) {
		pHemi = 1.0;
		pHomo = 0.0;
	    }

	    
	}


	if (Double.isNaN(L1) && Double.isNaN(L2)) {
	    pHemi = 1.0;
	    pHomo = 0.0;

	}


	// System.out.println("  muAB = " + muAB);

	Double fractionTumor = Double.NaN;
	if (pHomo <= pThreshold) {   // Hemi-deletion
	    fractionTumor = 2 - muAB;   
	}
	if (pHemi <= pThreshold) {  // Homo-deletion
	    fractionTumor = 1 - muAB / 2;
	}

	Double result[] = new Double[4];
	result[0] = pHemi;
	result[1] = pHomo;
	result[2] = fractionTumor;
	result[3] = new Double(segSize);
	return result;
    }



    public double getNormalTissueFraction() {
	return alpha;
    }

    public String getDeletionSegChrID(int i) {
	if (i >= 1 && i <= deletionSegChrID.size()) {
	    return deletionSegChrID.get(i - 1);
	} else {
	    return "Error!";
	}
    }


    public int getDeletionSegStartLocation(int i) {
	if (i >= 1 && i <= deletionSegStartLocation.size()) {
	    return deletionSegStartLocation.get(i - 1);
	} else {
	    return Integer.MIN_VALUE;
	}
    }


    public int getDeletionSegEndLocation(int i) {
	if (i >= 1 && i <= deletionSegEndLocation.size()) {
	    return deletionSegEndLocation.get(i - 1);
	} else {
	    return Integer.MIN_VALUE;
	}
    }

    public double getDeletionSegProbHemiDeletion(int i) {
	if (i >= 1 && i <= deletionSegProbHemiDeletion.size()) {
	    return deletionSegProbHemiDeletion.get(i - 1);
	} else {
	    return -Double.MAX_VALUE;
	}
    }

    public double getDeletionSegProbHomoDeletion(int i) {
	if (i >= 1 && i <= deletionSegProbHomoDeletion.size()) {
	    return deletionSegProbHomoDeletion.get(i - 1);
	} else {
	    return -Double.MAX_VALUE;
	}
    }


    public double getDeletionSegNormalTissueFraction(int i) {
	if (i >= 1 && i <= deletionSegNormalTissueFraction.size()) {
	    return deletionSegNormalTissueFraction.get(i - 1);
	} else {
	    return -Double.MAX_VALUE;
	}
    }



}





